ReduceFusion

对输入张量沿指定轴（axes）做归约（Reduce）运算，将若干维度压缩为标量结果，得到更低维的输出张量。 支持一次指定多个归约轴：按 axes 顺序依次规约，前一次的中间结果作为后一次的输入。

设待规约轴长度为 \(N\) （即 axis_size），该轴上元素为 \(x_0, x_1, \ldots, x_{N-1}\)， 则各 mode_ 对应计算如下：

\[\begin{split}\begin{aligned} \text{ReduceMean:} &\quad y = \frac{1}{N}\sum_{i=0}^{N-1} x_i \\ \text{ReduceMax:} &\quad y = \max_{i} x_i \\ \text{ReduceMin:} &\quad y = \min_{i} x_i \\ \text{ReduceProd:} &\quad y = \prod_{i=0}^{N-1} x_i \\ \text{ReduceSum:} &\quad y = \sum_{i=0}^{N-1} x_i \\ \text{ReduceSumSquare:} &\quad y = \sum_{i=0}^{N-1} x_i^{2} \\ \text{ReduceASum:} &\quad y = \sum_{i=0}^{N-1} |x_i| \\ \text{ReduceL2Norm:} &\quad y = \sqrt{\sum_{i=0}^{N-1} x_i^{2}} \end{aligned}\end{split}\]

计算过程：

1. 由输入形状与 axes 计算每个规约轴对应的 outer_size、axis_size、inner_size。

2. 将输入按 outer × axis × inner 展开；对每个 (outer, inner) 位置，沿 axis 维执行上述归约。

3. 若有多个规约轴，则重复步骤 2，中间结果写入 mid_buffer_0_ / mid_buffer_1_ 乒乓缓冲，最后一轮写到输出。

输入：

• src_data - 输入数据的地址

• param - 算子计算所需参数的结构体。其各成员见下述。

• core_mask - 核掩码（仅共享存储版本使用）。

ReduceParameter定义：

typedef struct ReduceParameter {
    void* tmp_src_data_;    // 单核 _p 临时输入缓冲，需 total_num_ * sizeof(T) 字节
    void* tmp_dst_data_;    // 单核 _p 临时输出缓冲，需 output_num_ * sizeof(T) 字节
    void* mid_buffer_0_;    // 多轴规约中间结果缓冲 0
    void* mid_buffer_1_;    // 多轴规约中间结果缓冲 1
    int* outer_sizes_;      // 各规约轴 outer size，需 num_axes_ * sizeof(int) 字节
    int* inner_sizes_;      // 各规约轴 inner size，需 num_axes_ * sizeof(int) 字节
    int* axis_sizes_;       // 各规约轴 axis size，需 num_axes_ * sizeof(int) 字节
    int total_num_;         // 输入元素总数
    int num_axes_;          // 规约轴数量
    int mode_;              // 规约模式
    int output_num_;        // 输出元素数
} ReduceParameter;

其中 T 为元素数据类型（如 fp32/int32 为 4 字节，fp16/int16 为 2 字节）。

缓冲分配说明：

• tmp_src_data_ / tmp_dst_data_：仅私有存储（_p）版本需要；共享存储（_s）可不设置。

• mid_buffer_0_ / mid_buffer_1_：各需 outer_sizes_[0] * inner_sizes_[0] * sizeof(T) 字节。

  • num_axes_ <= 1：两块均不使用

  • num_axes_ == 2：仅使用 mid_buffer_0_

  • num_axes_ >= 3：两块均需分配（乒乓复用）

• outer_sizes_ / inner_sizes_ / axis_sizes_：各需 num_axes_ 个 int。

• 若将 tmp_src_data_ / tmp_dst_data_ / mid_buffer_* 放在本核 AM，须避开算子硬编码工作区 [0x10000000, 0x10020400)，否则会被分块矩阵搬移覆盖。

其中 mode_ 的取值与规约方法的对应关系如下：

• ReduceMean = 0

• ReduceMax = 1

• ReduceMin = 2

• ReduceProd = 3

• ReduceSum = 4

• ReduceSumSquare = 5

• ReduceASum = 6

• ReduceL2Norm = 7

输出：

• dst_data - 输出地址。

支持平台：

FT78NE MT7004

备注

• FT78NE 支持 int8, int16, int32, fp32, fp64

• MT7004 支持 fp16, fp32, int16, int32

共享存储版本:

void i8_reduce_s(int8_t *src_data, int8_t *dst_data, ReduceParameter *param, int core_mask)

void i16_reduce_s(short *src_data, short *dst_data, ReduceParameter *param, int core_mask)

void i32_reduce_s(int *src_data, int *dst_data, ReduceParameter *param, int core_mask)

void hp_reduce_s(float16 *src_data, float16 *dst_data, ReduceParameter *param, int core_mask)

void fp_reduce_s(float *src_data, float *dst_data, ReduceParameter *param, int core_mask)

void dp_reduce_s(double *src_data, double *dst_data, ReduceParameter *param, int core_mask)

C调用示例：

// MT7004 示例（共享存储多核，DDR 地址）
void Resize(ReduceParameter* param, int ndim, int* input_shape, int num_axes, int* axes) {
    int tmp_input_shape[8];
    int total_num = 1;
    int i, j, k;
    for (i = 0; i < ndim; i++) {
        tmp_input_shape[i] = input_shape[i];
        total_num *= input_shape[i];
    }
    param->total_num_ = total_num;
    int offset_size = 0;
    for (i = 0; i < num_axes; ++i) {
        int axis = axes[i];
        int outer_size = 1;
        for (j = 0; j < axis; j++) {
            outer_size *= tmp_input_shape[j];
        }
        param->outer_sizes_[offset_size] = outer_size;
        int inner_size = 1;
        for (k = axis + 1; k < ndim; k++) {
            inner_size *= tmp_input_shape[k];
        }
        param->inner_sizes_[offset_size] = inner_size;
        param->axis_sizes_[offset_size] = tmp_input_shape[axis];
        offset_size++;
        tmp_input_shape[axis] = 1;
    }
    if (num_axes == 0) {
        param->output_num_ = total_num;
    } else {
        param->output_num_ = param->outer_sizes_[num_axes - 1] * param->inner_sizes_[num_axes - 1];
    }
}

void TestReduceSMCFp32(int* input_shape, int ndim, int* axes, int num_axes, int mode, int core_mask) {
    int core_id = get_core_id();
    int logic_core_id = GetLogicCoreId(core_mask, core_id);
    int core_num = GetCoreNum(core_mask);
    float* input = (float*)0x88000000;
    float* output = (float*)0x98000000;
    ReduceParameter* param = (ReduceParameter*)0xA8480000;
    if (logic_core_id == 0) {
        param->num_axes_ = num_axes;
        param->mode_ = mode;
        param->inner_sizes_ = (int*)0xA8484000;
        param->outer_sizes_ = (int*)0xA8485000;
        param->axis_sizes_ = (int*)0xA8486000;
        param->mid_buffer_0_ = (void*)0xA9490000;
        param->mid_buffer_1_ = (void*)0xAA490000;
        Resize(param, ndim, input_shape, num_axes, axes);
    }
    sys_bar(0, core_num); // 初始化参数完成后进行同步
    fp_reduce_s(input, output, param, core_mask);
}

void main(){
    int input_shape[3] = {4, 5, 5};
    int ndim = 3;
    int axes[1] = {1};
    int num_axes = 1;
    int mode = 0; // ReduceMean
    int core_mask = 0b1111;
    TestReduceSMCFp32(input_shape, ndim, axes, num_axes, mode, core_mask);
}

私有存储版本:

void i8_reduce_p(int8_t *src_data, int8_t *dst_data, ReduceParameter *param)

void i16_reduce_p(short *src_data, short *dst_data, ReduceParameter *param)

void i32_reduce_p(int *src_data, int *dst_data, ReduceParameter *param)

void hp_reduce_p(float16 *src_data, float16 *dst_data, ReduceParameter *param)

void fp_reduce_p(float *src_data, float *dst_data, ReduceParameter *param)

void dp_reduce_p(double *src_data, double *dst_data, ReduceParameter *param)

C调用示例：

// MT7004 示例（私有存储单核，AM 地址）
// tmp/mid 缓冲须避开 [0x10000000, 0x10020400)
void TestReduceAMFp32(int* input_shape, int ndim, int* axes, int num_axes, int mode) {
    float* input = (float*)0x10000000;
    float* output = (float*)0x10010000;
    float* tmp_src = (float*)0x10028000;
    float* tmp_dst = (float*)0x10038000;
    ReduceParameter* param = (ReduceParameter*)0x10048000;
    param->num_axes_ = num_axes;
    param->mode_ = mode;
    param->inner_sizes_ = (int*)0x10049000;
    param->outer_sizes_ = (int*)0x1004A000;
    param->axis_sizes_ = (int*)0x1004B000;
    param->mid_buffer_0_ = (void*)0x10050000;
    param->mid_buffer_1_ = (void*)0x10060000;
    param->tmp_src_data_ = tmp_src;
    param->tmp_dst_data_ = tmp_dst;
    Resize(param, ndim, input_shape, num_axes, axes);

    fp_reduce_p(input, output, param);
}

void main() {
    int input_shape[3] = {4, 5, 5};
    int ndim = 3;
    int axes[1] = {1};
    int num_axes = 1;
    int mode = 0; // ReduceMean
    TestReduceAMFp32(input_shape, ndim, axes, num_axes, mode);
}